import loguru
import time

from   accelerate import Accelerator
from   accelerate import load_checkpoint_in_model
from   accelerate import init_empty_weights
from   utils import get_loader, get_model

logger = loguru.logger
_, _, dataloader = get_loader()
model, optimizer, lr_scheduler = get_model()


model_directory = "model/accelerator.save_model"
accelerator     = Accelerator(gradient_accumulation_steps=4, mixed_precision="fp16")

#! 1. 模型的保存
accelerator.save_model(model=model, save_directory=model_directory, 
                        max_shard_size='200MB', safe_serialization=True)

#! 2. 模型加载
# with init_empty_weights():  #! 构建模型时不初始化参数
#     model, optimizer, lr_scheduler = get_model()
#     load_checkpoint_in_model(model, model_directory)
#     print(model.classifier.weight)